# src/step3_api_labeler.py
import os
import time
import requests  # 用于API调用
import json

# 导入项目配置和工具函数
from config import (
    DEEPSEEK_API_KEY, CANDIDATE_SENTENCES_FILE, LABELED_SENTENCES_API_FILE
)
from src.utils import setup_logger, load_json_file, save_json_file, get_progress_bar

# 为此模块设置独立的日志记录器
logger = setup_logger(__name__, "step3_api_labeling.log")

# DeepSeek API的端点 (请根据DeepSeek官方文档确认最新的API端点和请求格式)
# 这是一个假设的端点和模型名称，你需要替换为实际的
DEEPSEEK_API_URL = "https://api.deepseek.com/v1/chat/completions"  # 请查阅官方文档确认
DEEPSEEK_MODEL_NAME = "deepseek-chat"  # 或其他你希望使用的模型，例如 deepseek-coder (如果适用)

# 定义标签
LABEL_DESCRIPTION = "气候风险描述"
LABEL_COUNTERMEASURE = "气候风险应对措施"
VALID_LABELS = {LABEL_DESCRIPTION, LABEL_COUNTERMEASURE}


def call_deepseek_api(sentence, api_key, retries=3, delay=5):
    """
    调用DeepSeek API为单个句子打标签。

    Args:
        sentence (str): 需要打标签的句子。
        api_key (str): DeepSeek API密钥。
        retries (int): 失败后的重试次数。
        delay (int): 重试前的延迟秒数。

    Returns:
        str or None: 返回标签 ("气候风险描述" 或 "气候风险应对措施")，如果失败则返回None。
    """
    if not api_key or api_key == "YOUR_DEEPSEEK_API_KEY":
        logger.error("DeepSeek API Key 未配置或使用的是占位符。请在 config.py 中设置正确的API Key。")
        return None

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }

    # 构建prompt，引导模型进行二分类
    # 你可以根据API的实际效果调整这个prompt
    prompt_message = (
        f"请判断以下句子描述的是“气候风险描述”还是“气候风险应对措施”。"
        f"请只回答“{LABEL_DESCRIPTION}”或“{LABEL_COUNTERMEASURE}”。\n\n"
        f"句子：{sentence}\n\n"
        f"分类："
    )

    payload = {
        "model": DEEPSEEK_MODEL_NAME,
        "messages": [
            {"role": "system", "content": "你是一个专业的文本分类助手，专门识别与气候风险相关的句子类别。"},
            {"role": "user", "content": prompt_message}
        ],
        "max_tokens": 20,  # 限制输出长度，因为我们只需要一个短标签
        "temperature": 0.1  # 较低的temperature使输出更具确定性
    }

    for attempt in range(retries):
        try:
            logger.debug(f"向DeepSeek API发送请求 (尝试 {attempt + 1}/{retries})。句子: {sentence[:50]}...")
            response = requests.post(DEEPSEEK_API_URL, headers=headers, json=payload, timeout=30)  # 设置30秒超时
            response.raise_for_status()  # 如果HTTP状态码是4xx或5xx，则抛出异常

            response_data = response.json()
            logger.debug(f"API原始响应: {response_data}")

            if response_data.get("choices") and len(response_data["choices"]) > 0:
                # 提取模型返回的内容
                # 假设返回格式是: response_data['choices'][0]['message']['content']
                label_text = response_data["choices"][0].get("message", {}).get("content", "").strip()

                # 清洗和验证标签
                # 模型可能会返回一些额外文本，我们尝试提取核心标签
                if LABEL_DESCRIPTION in label_text:
                    return LABEL_DESCRIPTION
                elif LABEL_COUNTERMEASURE in label_text:
                    return LABEL_COUNTERMEASURE
                else:
                    logger.warning(f"API返回了无法识别的标签: '{label_text}' 对于句子: '{sentence}'")
                    # 可以考虑将原始返回也记录下来，或者尝试更复杂的解析
                    return None  # 或者返回原始文本供人工检查
            else:
                logger.warning(f"API响应中未找到有效的choices。响应: {response_data}")
                return None

        except requests.exceptions.HTTPError as http_err:
            logger.error(f"HTTP error occurred: {http_err} - {response.text}")
            if response.status_code == 401:  # Unauthorized
                logger.error("API认证失败，请检查你的DeepSeek API Key。")
                return None  # 认证错误不应重试
            if response.status_code == 429:  # Rate limit
                logger.warning(f"达到API速率限制，将在 {delay * (attempt + 2)} 秒后重试...")
                time.sleep(delay * (attempt + 2))  # 增加等待时间
            elif 500 <= response.status_code < 600:  # Server-side errors
                logger.warning(f"服务器端错误 ({response.status_code})，将在 {delay} 秒后重试...")
                time.sleep(delay)
            else:  # 其他客户端错误，可能不需要重试
                time.sleep(delay)  # 短暂等待
        except requests.exceptions.RequestException as req_err:
            logger.error(f"Request exception occurred: {req_err}")
            time.sleep(delay)
        except Exception as e:
            logger.error(f"调用DeepSeek API时发生未知错误: {e}")
            time.sleep(delay)

    logger.error(f"API调用失败 {retries} 次，放弃句子: '{sentence}'")
    return None


def run_api_labeling_step():
    """
    执行API打标签步骤的主函数。
    """
    logger.info("开始执行步骤3：API打标签...")

    # 1. 加载候选句子
    logger.info(f"加载候选句子从: {CANDIDATE_SENTENCES_FILE}")
    candidate_sentences = load_json_file(CANDIDATE_SENTENCES_FILE)

    if candidate_sentences is None:  # load_json_file 在失败时返回 None
        logger.error(f"未能加载候选句子文件: {CANDIDATE_SENTENCES_FILE}。请确保步骤2已成功运行。")
        return
    if not isinstance(candidate_sentences, list) or not all(isinstance(s, str) for s in candidate_sentences):
        logger.error(f"候选句子文件 {CANDIDATE_SENTENCES_FILE} 格式不正确，应为字符串列表。")
        return
    if not candidate_sentences:
        logger.warning("候选句子列表为空，无需进行API打标签。")
        save_json_file([], LABELED_SENTENCES_API_FILE)  # 保存一个空列表
        return

    logger.info(f"共加载 {len(candidate_sentences)} 个候选句子进行打标签。")

    labeled_data = []
    api_call_count = 0
    successful_labels = 0
    failed_labels = 0

    # 使用tqdm显示进度
    pbar = get_progress_bar(candidate_sentences, desc="Labeling sentences via API")

    for sentence in pbar:
        pbar.set_postfix_str(f"API Calls: {api_call_count}, Success: {successful_labels}, Fail: {failed_labels}",
                             refresh=True)

        label = call_deepseek_api(sentence, DEEPSEEK_API_KEY)
        api_call_count += 1

        if label and label in VALID_LABELS:
            labeled_data.append({"sentence": sentence, "label": label})
            successful_labels += 1
        else:
            # 即使API调用失败或返回无效标签，也记录下来，方便后续处理或人工校对
            labeled_data.append({"sentence": sentence, "label": "API_FAILED_OR_INVALID_LABEL"})
            failed_labels += 1
            logger.warning(f"句子未能成功打标或标签无效: '{sentence[:100]}...'")

        # API调用之间增加短暂延时，防止过于频繁请求 (根据API限制调整)
        time.sleep(0.5)  # 0.5秒延时，如果API有严格的QPS限制，需要增大

    logger.info("所有句子处理完毕。")
    logger.info(f"总API调用次数: {api_call_count}")
    logger.info(f"成功打标句子数: {successful_labels}")
    logger.info(f"失败或无效标签数: {failed_labels}")

    # 保存打好标签的数据
    logger.info(f"将打好标签的句子保存到: {LABELED_SENTENCES_API_FILE}")
    save_json_file(labeled_data, LABELED_SENTENCES_API_FILE)

    logger.info("步骤3：API打标签全部完成。")


if __name__ == '__main__':
    logger.info("独立运行API打标签模块测试...")

    # 1. 准备测试用的候选句子文件 (如果不存在则创建)
    # 这个文件通常由步骤2生成
    if not os.path.exists(CANDIDATE_SENTENCES_FILE):
        logger.warning(f"测试用的候选句子文件 {CANDIDATE_SENTENCES_FILE} 不存在，将创建示例文件。")
        sample_sentences = [
            "气候变化对农业生产构成显著风险。",
            "公司已采取措施减少碳排放，以应对气候风险。",
            "极端天气事件频发，导致保险赔付增加，这是重要的气候风险描述。",
            "我们投资于绿色技术，作为气候风险应对措施的一部分。",
            "海平面上升威胁沿海城市，这是一个长期的气候风险描述。",
            "制定应急预案以处理气候变化引发的自然灾害，是关键的应对措施。"
        ]
        save_json_file(sample_sentences, CANDIDATE_SENTENCES_FILE)

    # 检查API Key是否已配置
    if not DEEPSEEK_API_KEY or DEEPSEEK_API_KEY == "YOUR_DEEPSEEK_API_KEY":
        logger.error("=" * 50)
        logger.error("测试无法继续：DeepSeek API Key 未在 config.py 中配置。")
        logger.error("请将 'YOUR_DEEPSEEK_API_KEY' 替换为你的真实API密钥。")
        logger.error("=" * 50)
    else:
        logger.info(f"使用API Key: ...{DEEPSEEK_API_KEY[-4:]}")  # 显示部分API Key用于确认
        run_api_labeling_step()
        logger.info(f"API打标签模块独立测试运行结束。请检查输出文件: {LABELED_SENTENCES_API_FILE}")

